/**
 * 视频编码器实现
 * 主要功能：
 * 1. ARGB到YUV420SP格式转换
 * 2. 图像数据编码
 * 
 * @author sandyz987
 * @date 2023/11/17
 */

#include <cstdlib>
#include <cstdio>
#include <cmath>
#include <cstring>
#include <jni.h>
#include <map>
#include "utils.h"

// 1. 全局变量
jbyte yuv420spRaw[4096 * 2048];  // YUV420SP格式缓冲区

/**
 * 2. JNI接口函数
 * 将ARGB格式图像数据转换为YUV420SP格式
 * 
 * @param env JNI环境指针
 * @param jobj Java类对象
 * @param argb ARGB格式图像数据
 * @param width 图像宽度
 * @param height 图像高度
 * @return jbyteArray YUV420SP格式图像数据
 */
extern "C" jbyteArray
Java_com_sandyz_virtualcam_jni_EncoderJNI_encodeYUV420SP(JNIEnv *env, jclass jobj,
                                                         jintArray argb, jint width, jint height) {
    // 2.1 初始化变量
    int frameSize = width * height;  // 帧大小
    int yIndex = 0;                  // Y分量索引
    int uvIndex = frameSize;         // UV分量索引
    int a, R, G, B, Y, U, V;         // 颜色分量
    int index = 0;                   // 像素索引

    // 2.2 获取ARGB数据
    jint *argbArray = env->GetIntArrayElements(argb, 0);

    // 2.3 遍历图像像素
    for (int i = 0; i < height; ++i) {
        for (int j = 0; j < width; ++j) {
            // 2.3.1 提取ARGB分量
            a = argbArray[index] & -0x1000000 >> 24;  // Alpha通道
            R = argbArray[index] & 0xff0000 >> 16;    // 红色分量
            G = argbArray[index] & 0xff00 >> 8;       // 绿色分量
            B = argbArray[index] & 0xff;              // 蓝色分量

            // 2.3.2 RGB转YUV
            Y = ((66 * R + 129 * G + 25 * B + 128) >> 8) + 16;    // 亮度
            U = ((-38 * R - 74 * G + 112 * B + 128) >> 8) + 128;  // 色度U
            V = ((112 * R - 94 * G - 18 * B + 128) >> 8) + 128;   // 色度V

            // 2.3.3 存储Y分量
            yuv420spRaw[yIndex++] = (jbyte) (Y < 0 ? 0 : (Y > 255 ? 255 : Y));

            // 2.3.4 存储UV分量（每2x2像素采样一次）
            if (j % 2 == 0 && index % 2 == 0) {
                yuv420spRaw[uvIndex++] = (jbyte) (V < 0 ? 0 : (V > 255 ? 255 : V));
                yuv420spRaw[uvIndex++] = (jbyte) (U < 0 ? 0 : (U > 255 ? 255 : U));
            }
            index++;
        }
    }

    // 2.4 创建并返回YUV420SP数据
    jbyteArray yuv420sp = env->NewByteArray(frameSize * 3 / 2);
    env->SetByteArrayRegion(yuv420sp, 0, frameSize * 3 / 2, yuv420spRaw);
    return yuv420sp;
}